1141G - Privatization of Roads in Treeland - CodeForces Solution

#include<iostream>
#include<vector>
#include<map>
#include<set>
#include<cmath>
#include<cstring>
#include<algorithm>
#define int long long int
#define vec vector<int>
#define endl '\n'
#define f(i,n) for(int i=0;i<n;i++)
#define mod 1000000007
using namespace std;

int binpow(int a, int b) {
    a %= mod;
    int res = 1;
    while (b > 0) {
        if (b & 1)
            res = res * a % mod;
        a = a * a % mod;
        b >>= 1;
    }
    return res;
}

int add(int a, int b){
    a += b;
    if(a >= mod) a -= mod;
    return a;
}

int sub(int a, int b){
    a -= b;
    if(a < 0) a += mod;
    return a;
}

int mul(int a, int b){
    return (a%mod * b%mod)%mod;
}

map<int, vector<int>>adj;
map<pair<int,int>, int>e;
map<int,int>col;

void color(int i, int p, int prevcol, int m){
    int deg = adj[i].size();
    if(deg > m){
        for(int j:adj[i]){
            if(j == p) continue;
            col[e[{i,j}]] = 1;
            color(j,i,1,m);
        }
    }else{
        int curr = 1;
        for(int j:adj[i]){
            if(j == p) continue;
            if(prevcol == curr) curr++;
            col[e[{i,j}]] = curr;
            curr++;
            color(j,i,curr-1,m);
        }
    }
    
}
void sol(){
    int n; cin>>n;
    int k; cin>>k;
    f(i,n-1){
        int a,b; cin>>a>>b;
        adj[a].push_back(b);
        adj[b].push_back(a);
        e[{a,b}] = e[{b,a}] = i+1;
    }
    vec a;
    for(int i=1;i<=n;i++) a.push_back(adj[i].size());
    sort(a.begin(), a.end());
    int l = 1, r = n, ans = -1;
    while(l<=r){
        int m = (l+r)/2;
        auto it = upper_bound(a.begin(), a.end(),m);
        if(it == a.begin() || (n - (it - a.begin())) > k){
            l = m + 1;
        }else{
            ans = m;
            r = m - 1;
        }
    }
    cout<<ans<<endl;
    color(1,0,0,ans);
    for(int i=1;i<n;i++) cout<<col[i]<<" ";
}

signed main(){
    ios_base ::sync_with_stdio(0); 
    cin.tie(0);                    
    cout.tie(0);
    int t=1;
    //cin>>t;
    while(t--) sol();
}